# coding: utf-8

import os
import torch
from torch.utils.data import Dataset
import tiktoken
import json

class BuildDataset(Dataset):
    def __init__(self, path, block_size=512):
        super().__init__()
        # 采用 tiktoken 加载在线tokenizer
        self.enc = tiktoken.get_encoding("gpt2")
        
        self.block_size = block_size  # pos 最大长度
        self.encoded_data = []
        # 特殊符号分割文本， <|endoftext|>
        self.eos_token = self.enc.encode(
            "<|endoftext|>",
            allowed_special={"<|endoftext|>"}   
        )[0]
        
        raw_data = []
        self.max_lines = 1000
        with open(path, "r", encoding="utf-8") as f:
            for i, line in enumerate(f):
                data = json.loads(line.strip())
                if i >= self.max_lines:
                   break
                try:
                    text = data["text"]
                    raw_data.append(text)
                except Exception as e:
                    continue
        
        full_encoded = []
        for text in raw_data:
            encoded_text = self.enc.encode(text)
            full_encoded.extend(encoded_text + [self.eos_token])
            
        # block_size 是 512
        for i in range(0, len(full_encoded), self.block_size):
            chunk = full_encoded[i:i + self.block_size + 1]  # 512 shift 1 位 每一行实际长度是513个token
            if len(chunk) < self.block_size + 1:
                chunk = chunk + [self.eos_token] * (self.block_size + 1 - len(chunk))
            self.encoded_data.append(chunk)
        
    def __len__(self):
        return len(self.encoded_data)

    def __getitem__(self, idx):
        # 在处理数据时，就进行文本的shift操作，输入是 前512个token，输出是 后512个token
        chunk = self.encoded_data[idx]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y

    def encode(self, text):
        # 将文本编码为token ids
        return self.enc.encode(text)
    
    def decode(self, ids):
        # 将token ids解码为文本
        return self.enc.decode(ids)